import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ------------------------------
# 传感器数据
# ------------------------------
top_data = {
    "In0_1": 23.84770966, "In0_2": -6.376930237, "In0_3": -0.194458008, "Mid0": 1.028503418,
    "In45_1": -3.536563873, "In45_2": -0.87474823, "In45_3": 0.700706482, "Mid45": -1.003303528,
    "In90_1": -0.985961914, "In90_2": -0.574142456, "In90_3": -4.267082214, "Mid90": -1.469314575,
    "In135_1": 0.387619019, "In135_2": -0.473937988, "In135_3": 1.510803223, "Mid135": 0.003700256,
    "In180_1": 24.58776855, "In180_2": -1.974632263, "In180_3": -0.345870972, "Mid180": -0.363586426,
    "In225_1": -1.766098022, "In225_2": 0.067398071, "In225_3": -0.177062988, "Mid225": -3.872184753,
    "In270_1": 0.21295166, "In270_2": -3.948898315, "In270_3": -2.954559326, "Mid270": -1.540847778,
    "In315_1": -0.123123169, "In315_2": -3.276229858, "In315_3": 0.038719177, "Mid315": 0.487678528
}

bottom_data = {
    "In0_1": 570.4158325, "In0_2": 550.4851685, "In0_3": 561.7962646, "Mid0": 589.0463257,
    "In45_1": 529.7352905, "In45_2": 536.1941528, "In45_3": 540.2058105, "Mid45": 551.6915894,
    "In90_1": 512.1013184, "In90_2": 509.1639404, "In90_3": 508.4290771, "Mid90": 494.1642456,
    "In135_1": 499.0268555, "In135_2": 491.5413513, "In135_3": 484.3640137, "Mid135": 459.7319336,
    "In180_1": 517.0548096, "In180_2": 490.5742798, "In180_3": 484.4238586, "Mid180": 458.5591431,
    "In225_1": 512.7064819, "In225_2": 511.1669922, "In225_3": 507.770874, "Mid225": 501.3547974,
    "In270_1": 532.8231201, "In270_2": 535.7677612, "In270_3": 539.0004883, "Mid270": 544.1796875,
    "In315_1": 548.0007324, "In315_2": 554.7557373, "In315_3": 564.3328247, "Mid315": 591.8924561
}

adjusted_data = {
    "In0_1": 357.5815125, "In0_2": 332.3037109, "In0_3": 341.5246277, "Mid0": 356.9660645,
    "In45_1": 323.3222656, "In45_2": 328.7756042, "In45_3": 331.3685913, "Mid45": 338.0082397,
    "In90_1": 316.5126038, "In90_2": 315.7456665, "In90_3": 315.415741, "Mid90": 308.4900513,
    "In135_1": 311.2686157, "In135_2": 307.0535583, "In135_3": 303.7873535, "Mid135": 291.7260132,
    "In180_1": 318.3375854, "In180_2": 305.3764648, "In180_3": 303.2308044, "Mid180": 290.2286987,
    "In225_1": 316.5319824, "In225_2": 316.7035828, "In225_3": 315.1881714, "Mid225": 312.3748779,
    "In270_1": 307.3053894, "In270_2": 327.0389404, "In270_3": 328.885498, "Mid270": 327.5952454,
    "In315_1": 335.2342834, "In315_2": 336.7516785, "In315_3": 343.2346802, "Mid315": 357.8313599
}

# ------------------------------
# 传感器坐标生成
# ------------------------------
def generate_sensor_positions():
    angles = np.deg2rad([0, 45, 90, 135, 180, 225, 270, 315])
    radii = [45000, 60000, 72500, 120000]  # 内圈到外圈的半径比例
    positions = {}
    
    for angle in angles:
        deg = int(np.rad2deg(angle))
        for j, r in enumerate(radii):
            x = r * np.cos(angle)
            y = r * np.sin(angle)
            positions[f"In{deg}_{j+1}"] = (x, y)
            if j == 3:  # 最外圈的传感器标记为Mid
                positions[f"Mid{deg}"] = (x, y)
    return positions

sensor_pos = generate_sensor_positions()

# ------------------------------
# 平面拟合函数
# ------------------------------
def fit_plane(data, sensor_pos):
    """使用最小二乘法拟合平面 z = a*x + b*y + c"""
    points = []
    for name, value in data.items():
        if name in sensor_pos:
            x, y = sensor_pos[name]
            points.append([x, y, value])
    
    if not points:
        raise ValueError("没有找到匹配的传感器数据")
    
    points = np.array(points)
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    
    A = np.column_stack([x, y, np.ones_like(x)])
    coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
    a, b, c = coeffs
    
    residuals = z - (a*x + b*y + c)
    rms_error = np.sqrt(np.mean(residuals**2))
    
    return a, b, c, rms_error

# ------------------------------
# 计算平面参数
# ------------------------------
try:
    a_top, b_top, c_top, top_error = fit_plane(top_data, sensor_pos)
    a_bottom, b_bottom, c_bottom, bottom_error = fit_plane(bottom_data, sensor_pos)
    a_adjusted, b_adjusted, c_adjusted, adjusted_error = fit_plane(adjusted_data, sensor_pos)
    
    print("Topchuck 平面拟合结果:")
    print(f"方程: z = {a_top:.4f}x + {b_top:.4f}y + {c_top:.4f}")
    print(f"RMS误差: {top_error:.4f}\n")
    
    print("Bottomchuck 原始平面拟合结果:")
    print(f"方程: z = {a_bottom:.4f}x + {b_bottom:.4f}y + {c_bottom:.4f}")
    print(f"RMS误差: {bottom_error:.4f}\n")
    
    print("Bottomchuck 调平后平面拟合结果:")
    print(f"方程: z = {a_adjusted:.4f}x + {b_adjusted:.4f}y + {c_adjusted:.4f}")
    print(f"RMS误差: {adjusted_error:.4f}\n")

except ValueError as e:
    print(f"错误: {e}")
    exit()

# ------------------------------
# 可视化函数
# ------------------------------
def plot_plane_and_data(ax, data, sensor_pos, a, b, c, title, color='blue'):
    """绘制传感器数据和拟合平面"""
    points = []
    for name, value in data.items():
        if name in sensor_pos:
            x, y = sensor_pos[name]
            points.append([x, y, value])
    
    points = np.array(points)
    x_data = points[:, 0]
    y_data = points[:, 1]
    z_data = points[:, 2]
    
    # 绘制数据点
    ax.scatter(x_data, y_data, z_data, c=color, s=50, depthshade=True, label=f'{title} Data')
    
    # 生成拟合平面
    x_min, x_max = min(x_data)-1, max(x_data)+1
    y_min, y_max = min(y_data)-1, max(y_data)+1
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 20),
                         np.linspace(y_min, y_max, 20))
    zz = a * xx + b * yy + c
    
    # 绘制平面
    ax.plot_surface(xx, yy, zz, alpha=0.5, color=color)
    
    ax.set_xlabel('X Position (μm)')
    ax.set_ylabel('Y Position (μm)')
    ax.set_zlabel('Z Value')
    ax.set_title(title)

# ------------------------------
# 创建可视化
# ------------------------------
fig = plt.figure(figsize=(18, 6))

# Topchuck 可视化
ax1 = fig.add_subplot(131, projection='3d')
plot_plane_and_data(ax1, top_data, sensor_pos, a_top, b_top, c_top, 'Topchuck', 'blue')

# Bottomchuck 原始可视化
ax2 = fig.add_subplot(132, projection='3d')
plot_plane_and_data(ax2, bottom_data, sensor_pos, a_bottom, b_bottom, c_bottom, 'Original Bottom', 'red')

# Bottomchuck 调平后可视化
ax3 = fig.add_subplot(133, projection='3d')
plot_plane_and_data(ax3, adjusted_data, sensor_pos, a_adjusted, b_adjusted, c_adjusted, 'Adjusted Bottom', 'green')

plt.tight_layout()

# ------------------------------
# 计算调平参数
# ------------------------------
def calculate_adjustment(a_top, b_top, a_bottom, b_bottom):
    """计算需要调整的rx和ry（微弧度）"""
    rx_rad = np.arctan(a_bottom) - np.arctan(a_top)
    ry_rad = np.arctan(b_bottom) - np.arctan(b_top)
    return rx_rad * 1e6, ry_rad * 1e6

rx_µrad, ry_µrad = calculate_adjustment(a_top, b_top, a_bottom, b_bottom)
adjusted_rx_µrad, adjusted_ry_µrad = calculate_adjustment(a_top, b_top, a_adjusted, b_adjusted)

print("\n原始调平参数计算结果:")
print(f"需要调整的 rx: {rx_µrad:.2f} µrad")
print(f"需要调整的 ry: {ry_µrad:.2f} µrad")

print("\n实际调整后的参数:")
print(f"调整后的 rx: {adjusted_rx_µrad:.2f} µrad")
print(f"调整后的 ry: {adjusted_ry_µrad:.2f} µrad")

print("\n调整效果:")
print(f"rx 调整量: {rx_µrad-adjusted_rx_µrad:.2f} µrad")
print(f"ry 调整量: {ry_µrad-adjusted_ry_µrad:.2f} µrad")

plt.show()